package cn.syx.service;

import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentWriter;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.transformer.splitter.TextSplitter;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * rag
 * DocumentReader  加载文本内容 返回 Document
 * DocumentTransformer
 */
@Component
public class RagService {

    @Value("classpath:meituan-qa.txt")
    private Resource resource;


    /**
     * 加载文件资源 获取 Document
     * @return
     */
    public List<Document> loadText(){
        TextReader textReader = new TextReader(resource);
        textReader.getCustomMetadata().put("filename","meituan-qa.txt");
        List<Document> documents = textReader.get();
        return documents;
    }

    /**
     * 将 Document 以 token进行切分，获得切分后的 Document
     * @param documents
     * @return
     */
    public List<Document> tokenTextSplitter(List<Document> documents){

        TokenTextSplitter tokenTextSplitter = new TokenTextSplitter();
        List<Document> documentList = tokenTextSplitter.apply(documents);
        return documentList;
    }

    /**
     * 使用自定义切分器  SyxTextSplitter
     * @param documents
     * @return
     */
    public List<Document> syxTextSplitter(List<Document> documents){

        SyxTextSplitter syxTextSplitter = new SyxTextSplitter();
        List<Document> documentList = syxTextSplitter.apply(documents);
        return documentList;
    }

    // DocumentWriter 向量存储内存（接口）    ，有个子接口 VectorStore  的默认实现 SimpleVectorStore
    //所以需要注入 VectorStore 接口bean
    @Autowired
    private VectorStore vectorStore;

    /**
     * 向量存储
     * @param documents
     */
    public void documentWriter(List<Document> documents){
        //将切分后的 Document 向量存入 内存中
        vectorStore.add(documents);
    }

    /**
     * 向量查找
     * @param message
     * @return
     */
    public List<Document> search(String message){
        List<Document> documents = vectorStore.similaritySearch(message);
        return documents;
    }

    // 使用 按 token进行切分 ，无法获得 一个问题一个答案的模式，需要自定义切分器

    public class SyxTextSplitter extends TextSplitter {

        @Override
        protected List<String> splitText(String text) {
            return List.of(split(text));
        }

        public String[] split(String text){
            return text.split("\\s*\\R\\s*\\R\\s*");
        }
    }


}
